#!/usr/bin/env python3

import argparse
import f90nml
import pendulum
from glob import glob
from collections import OrderedDict
import re
import os
from pprint import pprint

def parse_time(string):
	match = re.match(r'(\d{4}\d{2}\d{2}\d{2})(\d{2})?', string)
	if match.group(2):
		return pendulum.from_format(string, 'YYYYMMDDHHmm')
	else:
		return pendulum.from_format(string, 'YYYYMMDDHH')

parser = argparse.ArgumentParser(description='Write namelist for EFSO run.')
parser.add_argument('--time', help='Forecast start time (YYYYMMDDHH)', type=parse_time)
parser.add_argument('--use-synop', action='store_true')
parser.add_argument('--synop-root', default='/data/raw/cimiss/SURF_CHN_MAIN_MIN')
parser.add_argument('--synop-file-pattern', default='YYYYMMDDHHmm')
parser.add_argument('--use-raob', action='store_true')
parser.add_argument('--raob-root', default='/data/raw/cimiss/UPAR_CHN_MUL_FTM')
parser.add_argument('--raob-file-pattern', default='YYYYMMDDHHmm')
parser.add_argument('--use-amdar', action='store_true')
parser.add_argument('--amdar-root', default='/data/raw/cimiss/UPAR_ARD_G_MUL_MUT_TAB')
parser.add_argument('--amdar-file-pattern', default='YYYYMMDDHHmm')
parser.add_argument('--use-profiler', action='store_true')
parser.add_argument('--profiler-root', default='/data/raw/profiler')
parser.add_argument('--grapes-root', default='/data/raw/grapes_geps')
parser.add_argument('--fcst-hour', default=24, type=int)
parser.add_argument('--era5-root', default='/data/raw/era5')
args = parser.parse_args()

namelist = {
	'efso': OrderedDict([
		('use_synop', False),
		('synop_files', []),
		('use_raob', False),
		('raob_files', []),
		('use_amdar', False),
		('amdar_files', []),
		('use_profiler', False),
		('profiler_files', []),
		('model_files_t0', []),
		('model_files_tn', []),
		('era5_file_tn', ''),
		('ensemble_size', 31),
		('parallel_ensemble', True),
		('obs_impact_file', '')
	])
}

namelist['efso']['use_synop'] = args.use_synop
namelist['efso']['synop_files'] = [f'{args.synop_root}/{args.time.format(args.synop_file_pattern)}.xml']

namelist['efso']['use_raob'] = args.use_raob
namelist['efso']['raob_files'] = [f'{args.raob_root}/{args.time.format(args.raob_file_pattern)}.xml']

namelist['efso']['use_amdar'] = args.use_amdar
namelist['efso']['amdar_files'] = [f'{args.amdar_root}/{args.time.format(args.amdar_file_pattern)}.xml']

namelist['efso']['use_profiler'] = args.use_profiler
# Generate a file list for profilers.


station_dir = os.listdir('{}/{}/{}/{}'.format(args.profiler_root,args.time.format('YYYY'),args.time.format('MM'),args.time.format('DD')))
conver_file = []
file_list   = []
for i in station_dir:
    file_list += glob('{}/{}/{}/{}/{}/ROBS/QI/Z_RADA*'.format(args.profiler_root,args.time.format('YYYY'),args.time.format('MM'),args.time.format('DD'), i))
# /nas02/data/raw/profiler/2021/06/18/51463/ROBS/QI/
for pro_file in file_list:
    time = pendulum.from_format(re.findall(r'\d{14}', os.path.basename(pro_file))[0], 'YYYYMMDDHHmmss')
    if time == args.time:
        conver_file += pro_file + '\n'
if len(conver_file) < 1:
    print(f'[Warning]{args.time.format("YYYYMMDDHH")} is empty')
    namelist['efso']['use_profiler'] = False
    namelist['efso']['profiler_files'] = 0
else:
    pro_list_file = f'/tmp/pro_list_{args.time.format("YYYYMMDDHH")}.txt'
    with open(pro_list_file, 'w') as f:
        f.writelines(conver_file)
    namelist['efso']['profiler_files'] = [f'{pro_list_file}']

for iens in range(31):
	namelist['efso']['model_files_t0'].append(f'{args.grapes_root}/{args.time.format("YYYYMMDDHH")}/gef.gra.{str(iens).zfill(3)}.{args.time.format("YYYYMMDDHH")}000.grb2')
	namelist['efso']['model_files_tn'].append(f'{args.grapes_root}/{args.time.format("YYYYMMDDHH")}/gef.gra.{str(iens).zfill(3)}.{args.time.format("YYYYMMDDHH")}{str(args.fcst_hour).zfill(3)}.grb2')

era5_time = args.time.add(hours=args.fcst_hour)
namelist['efso']['grapes_file_tn'] = f'{args.grapes_root}/{era5_time.format("YYYYMMDDHH")}/gef.gra.{str(0).zfill(3)}.{era5_time.format("YYYYMMDDHH")}000.grb2'
namelist['efso']['era5_file_tn'] = f'{args.era5_root}/{era5_time.format("YYYYMMDD")}/era5_{era5_time.format("YYYYMMDDHH")}.nc'

namelist['efso']['obs_impact_file'] = f'obs_impact.{args.time.format("YYYYMMDDHH")}.txt'

namelist['efso']['omb_ua_limit'] = 20
namelist['efso']['omb_va_limit'] = 20
namelist['efso']['omb_ta_limit'] = 10

f90nml.write(namelist, open('namelist', 'w'))
print('[Notice]: File namelist is created.')
